#include "MemoryManager.h"
#include "MemoryManagerCore.h"
#include "Constants.h"
#include "StackAllocator.h"
#include "HeapAllocator.h"
#include "PoolAllocator.h"
#include "Assertions.h"
#include "Console.h"

#include <cstring>
#include <new>

namespace CEngine
{

namespace MemoryManagerNothingToSeeHere {
    StackAllocator* tempStack;
    U8 currentTempStack;

    HeapAllocator heap(Constants::Memory::heapMemory, Constants::Memory::heapPageSize);
    U64 heapMin;
    U64 heapMax;

    PoolAllocator pool24(24, Constants::Memory::pool24Memory);
    U64 pool24Min;
    U64 pool24Max;

    StackAllocator stack(Constants::Memory::stackMemory);
    U64 stackMin;
    U64 stackMax;
    void stackAlloc(void* pointer, U64 bytes);
    void stackFree(void* pointer);

    bool initialized = false;

    void testMemory(void* pointer, U64 bytes, bool removeEntry);
}

namespace {
    struct MemoryRange {
        U64 beginning;
        U64 end;
    };
    struct StackHeader {
        void* pointer;
        U64 size;
    };

    MemoryRange* freedStack;
    U64 freedStackSize, freedStackCapacity;
    StackHeader* stackHeader;
    U64 stackHeaderSize, stackHeaderCapacity;

    bool growStackAlloc = false;

    MemoryRange* trackedMemory;
    U64 trackedAllocations;
}

namespace MemoryManager {

    using namespace MemoryManagerNothingToSeeHere;

    void
    initialize(void) {
        tempStack = (StackAllocator*)new U8[sizeof(StackAllocator) * Constants::Memory::tempFrames];
        for (U8 i = 0; i < Constants::Memory::tempFrames; ++i) {
            new(&tempStack[i]) StackAllocator(Constants::Memory::tempMemory);
        }
        currentTempStack = 0;

        heapMin = (U64)heap.heapPointer;
        heapMax = heapMin + heap.getTotalBytes();

        pool24Min = (U64)pool24.poolPointer;
        pool24Max = pool24Min + (pool24.getTotalElements() * 24);

        stackMin = (U64)stack.stackPointer;
        stackMax = stackMin + stack.totalBytes;

        if(Constants::Memory::trackMemory > 0) {
            trackedAllocations = 0;
            trackedMemory = new MemoryRange[655360];
        }

        freedStackSize = 0;
        freedStackCapacity = 1024;
        freedStack = new MemoryRange[freedStackCapacity];
        stackHeaderSize = 0;
        stackHeaderCapacity = 1024;
        stackHeader = new StackHeader[stackHeaderCapacity];

        initialized = true;

        defaultLog.log(Logging::MEMORY, LOG_INFO, "Memory manager initialized.\n");
    }

    void*
    tempAlloc(U32 bytes) {
        ASSERT(initialized);

        if(bytes == 0) {
            return NULL;
        }

        void* pointer = tempStack[currentTempStack].alloc(bytes);
        if(!pointer) {
            defaultLog.log(Logging::MEMORY, LOG_ERROR, "Temperary frame stack is full!\n");
        }
        return pointer;
    }

    void*
    rawAlloc(U64 bytes) {
        ASSERT(initialized);

        void* pointer;
        if(bytes == 0) {
            return NULL;
        }
        else if(bytes == 24) {
            pointer = pool24.alloc();

            static U8 errorLevel = 0;
            U64 elementsRemaining = pool24.getFreeElements();
            if(elementsRemaining == 0) {
                if(errorLevel < 2) {
                    defaultLog.log(Logging::MEMORY, LOG_ERROR, "24 byte pool empty!\n");
                }
                errorLevel = 2;
            }
            else if(elementsRemaining < (Constants::Memory::pool24Memory * 0.05)) {
                if(errorLevel < 1) {
                    defaultLog.log(Logging::MEMORY, LOG_WARNING, "Remaining 24 byte pool elements below 5%%!\n");
                }
                errorLevel = 1;
            }
            else {
                errorLevel = 0;
            }
        }
        else if(/*bytes < Constants::Memory::heapPageSize &&*/ bytes <= stack.getFreeBytes()) {
            pointer = stack.alloc(bytes);
            stackAlloc(pointer, bytes);

            static bool lowMemory = false;
            U64 bytesRemaining = stack.getFreeBytes();
            if(bytesRemaining < (Constants::Memory::stackMemory * 0.05)) {
                if(!lowMemory) {
                    defaultLog.log(Logging::MEMORY, LOG_WARNING, "Remaining stack memory below 5%%!\nMay begin to use heap memory instead!\n");
                }
                lowMemory = true;
            }
            else {
                lowMemory = false;
            }
        }
        else {
            pointer = heap.alloc(bytes);

            static U8 errorLevel = 0;
            U64 byteRemaining = heap.getFreeBytes();
            if(byteRemaining == 0) {
                if(errorLevel < 2) {
                    defaultLog.log(Logging::MEMORY, LOG_ERROR, "heap empty!\n");
                }
                errorLevel = 2;
            }
            else if(byteRemaining < (Constants::Memory::heapMemory * 0.05)) {
                if(errorLevel < 1) {
                    defaultLog.log(Logging::MEMORY, LOG_WARNING, "Remaining heap memory below 5%%!\n");
                }
                errorLevel = 1;
            }
            else {
                errorLevel = 0;
            }
        }

        if(!pointer) {
            float outputBytes = bytes;
            if(outputBytes > (10 * 1024 * 1024)) {
                outputBytes = outputBytes / (1024 * 1024);
                defaultLog.log(Logging::MEMORY, LOG_ERROR, "Failed to allocate %.2fMB!\n", outputBytes);
            }
            else if(outputBytes > (10 * 1024)) {
                outputBytes = outputBytes / (1024);
                defaultLog.log(Logging::MEMORY, LOG_ERROR, "Failed to allocate %.2fkB!\n", outputBytes);
            }
            else {
                defaultLog.log(Logging::MEMORY, LOG_ERROR, "Failed to allocate %.0fB!\n", outputBytes);
            }
        }
        testMemory(pointer, bytes, false);
        return pointer;
    }

    void*
    grow(void* pointer, U64 oldSize, U64 newSize) {
        ASSERT(oldSize <= newSize);
        ASSERT(initialized);

        if(!pointer) {
            return NULL;
        }

        if(oldSize == newSize) {
            return pointer;
        }

        U8 allocator;
        if((U64)pointer >= heapMin && (U64)pointer <= heapMax) {
            allocator = 1;
        }
        else if((U64)pointer >= pool24Min && (U64)pointer <= pool24Max) {
            allocator = 2;
        }
        else if((U64)pointer >= stackMin && (U64)pointer <= stackMax) {
            allocator = 3;
        }
        else {
            defaultLog.log(Logging::MEMORY, LOG_ERROR, "pointer was not from memory manager!\n");
            return pointer;
        }

        void* newPointer = pointer;
        U64 difference = newSize - oldSize;

        switch(allocator) {
        case 1/*heap*/:
            if(((oldSize % heap.getPageSize()) + difference) <= heap.getPageSize()) {
                break;
            }

            newPointer = rawAlloc(newSize);

            if(!newPointer) {
                break;
            }

            memcpy(newPointer, pointer, oldSize);

            free(pointer);
            break;
        case 2/*pool24*/:
            newPointer = rawAlloc(newSize);

            if(!newPointer) {
                break;
            }

            memcpy(newPointer, pointer, oldSize);

            free(pointer);
            break;
        case 3/*stack*/:
            if((StackAllocator::Marker)(pointer + oldSize) == stack.getMarker() && difference < heap.getPageSize() && difference < stack.getFreeBytes()) {
                if(difference == 24) {
                    ++difference;
                }
                growStackAlloc = true;
                rawAlloc(difference);
                growStackAlloc = false;
                break;
            }

            newPointer = rawAlloc(newSize);

            if(!newPointer) {
                break;
            }

            memcpy(newPointer, pointer, oldSize);

            free(pointer);
            break;
        }

        return newPointer;
    }

    void
    incrementTempStack(void) {
        ASSERT(initialized);
        ++currentTempStack;
        if(currentTempStack >= Constants::Memory::tempFrames) {
            currentTempStack = 0;
        }
        tempStack[currentTempStack].clear();
    }

    void
    printMemoryUsage(void) {
        if(!initialized) {
            Console::print("%qMemory manager not initialized!\n%r", Console::LRED);
            return;
        }

        U64 tempFreeSpace[Constants::Memory::tempFrames];
        for (U8 i = 0; i < Constants::Memory::tempFrames; ++i) {
            tempFreeSpace[i] = tempStack[i].getFreeBytes();
        }

        U64 stackFreeSpace = stack.getFreeBytes();
        U64 heapFreeSpace = heap.getFreeBytes();
        U64 pool24FreeSpace = pool24.getFreeElements();

        U64 stackPendingFreeSpace = 0;
        for (U64 it = 0; it < freedStackSize; ++it) {
            stackPendingFreeSpace += freedStack[it].end - freedStack[it].beginning;
        }

        U8 colour;

        Console::print("%q|MEMORY MANAGER|\n%r", Console::LRED);

        colour = Console::defaultConsoleTextColour;
        if(stackFreeSpace <= (stack.getTotalBytes() * 0.1f)) {
            colour = (((colour | (U8)255) - 255) + Console::LYELLOW);
        }
        if(stackFreeSpace <= (stack.getTotalBytes() * 0.01f)) {
            colour = (((colour | (U8)255) - 255) + Console::LRED);
        }
        Console::print("%qStack: %llu bytes free of %llu bytes (%llu bytes pending to be freed).\n\r", colour, stackFreeSpace, stack.getTotalBytes(), stackPendingFreeSpace);

        colour = Console::defaultConsoleTextColour;
        if(heapFreeSpace <= (heap.getTotalBytes() * 0.1f)) {
            colour = (((colour | (U8)255) - 255) + Console::LYELLOW);
        }
        if(heapFreeSpace <= (heap.getTotalBytes() * 0.01f)) {
            colour = (((colour | (U8)255) - 255) + Console::LRED);
        }
        Console::print("%qHeap: %llu bytes free of %llu bytes.\n%r", colour, heapFreeSpace, heap.getTotalBytes());

        colour = Console::defaultConsoleTextColour;
        if(pool24FreeSpace <= (pool24.getTotalElements() * 0.1f)) {
            colour = (((colour | (U8)255) - 255) + Console::LYELLOW);
        }
        if(pool24FreeSpace <= (pool24.getTotalElements() * 0.01f)) {
            colour = (((colour | (U8)255) - 255) + Console::LRED);
        }
        Console::print("%q24 Byte Pool: %llu elements free of %llu elements.\n%r", colour, pool24FreeSpace, pool24.getTotalElements());
        for (U64 i = 0; i < Constants::Memory::tempFrames; ++i) {
            colour = Console::defaultConsoleTextColour;
            if(tempFreeSpace[i] <= (tempStack[i].getTotalBytes() * 0.1f)) {
                colour = (((colour | (U8)255) - 255) + Console::LYELLOW);
            }
            if(tempFreeSpace[i] <= (tempStack[i].getTotalBytes() * 0.01f)) {
                colour = (((colour | (U8)255) - 255) + Console::LRED);
            }
            Console::print("%qTemp Stack %u: %llu bytes free of %llu bytes.\n%r", colour, i, tempFreeSpace[i], tempStack[i].getTotalBytes());
        }
        CHAR line[40];
        for (U8 i = 0; i < 40; ++i) {
            line[i] = '-';
        }
        line[39] = '\0';
        Console::print("%s\n", line);
    }

    void
    printStackUsage(void) {
        if(!initialized) {
            Console::print("%qMemory manager not initialized!\n%r", Console::LRED);
            return;
        }

        U32 resolution = 2;
        U32 blockSize = stack.getTotalBytes() / resolution;
        if((blockSize % resolution) != 0) {
            ++blockSize;
        }

        Logging tempLog(LOG_DEBUG, false, "stackLog.txt", 1100000);
        tempLog.log(Logging::MEMORY, LOG_INFO, "|Stack Memory|");

        CHAR block[blockSize];
        for (U32 i = 0; i < blockSize; ++i) {
            block[i] = ' ';
        }
        block[blockSize - 1] = '\0';

        U64 stackUsedBytes = stack.getTotalBytes() - stack.getFreeBytes();
        U32 temp = stackUsedBytes / resolution;
        if((stackUsedBytes % resolution) != 0) {
            ++temp;
        }

        for (U32 i = 0; i < temp; ++i) {
            block[i] = '-';
        }

        tempLog.log(Logging::MEMORY, LOG_INFO, "|%s|", block);

        for (U32 i = 0; i < blockSize; ++i) {
            block[i] = ' ';
        }
        block[blockSize - 1] = '\0';

        for (U32 i = 0; i < freedStackSize; ++i) {
            U64 beginBlock = (freedStack[i].beginning - (stack.getMarker() - stackUsedBytes)) / resolution;
            U64 endBlock = (freedStack[i].end - (stack.getMarker() - stackUsedBytes)) / resolution;

            for (U64 i = beginBlock; i <= endBlock; ++i) {
                block[i] = '-';
            }
        }

        tempLog.log(Logging::MEMORY, LOG_INFO, "|%s|", block);
    }

    U64
    getTempFreeBytes(void) {
        ASSERT(initialized);
        return tempStack[currentTempStack].getFreeBytes();
    }

    U64
    getHeapFreeBytes(void) {
        ASSERT(initialized);
        return heap.getFreeBytes();
    }

    U64
    getStackFreeBytes(void) {
        ASSERT(initialized);
        return stack.getFreeBytes();
    }

    U64
    get24PoolFreeElements(void) {
        ASSERT(initialized);
        return pool24.getFreeElements();
    }

    U64
    getHeapPageSize(void) {
        return Constants::Memory::heapPageSize;
    }

    void
    shutdown(void) {
        if(tempStack) {
            for (U8 i = 0; i < Constants::Memory::tempFrames; ++i) {
                tempStack[i].~StackAllocator();
            }
            delete[] (U8*)tempStack;
            tempStack = NULL;
        }

        heap.clear();
        pool24.clear();
        stack.clear();

        delete freedStack;
        delete stackHeader;

        if(Constants::Memory::trackMemory > 0) {
            trackedAllocations = 0;
            delete trackedMemory;
        }

        initialized = false;
        defaultLog.log(Logging::MEMORY, LOG_INFO, "Memory manager shutdown.\n");
    }
}

namespace MemoryManagerNothingToSeeHere {
    void
    stackAlloc(void* pointer, U64 bytes) {
        if(growStackAlloc) {
            stackHeader[stackHeaderSize - 1].size += bytes;
        }
        else {
            StackHeader header;
            header.pointer = pointer;
            header.size = bytes;

            if(stackHeaderSize == stackHeaderCapacity) {
                StackHeader* temp = stackHeader;
                stackHeaderCapacity *= 2;
                stackHeader = new StackHeader[stackHeaderCapacity];
                memcpy(stackHeader, temp, stackHeaderSize);
                delete temp;
            }

            stackHeader[stackHeaderSize++] = header;
        }
    }

    void
    stackFree(void* pointer) {
        StackHeader header;
        header.pointer = NULL;

        for (U64 it = 0; it < stackHeaderSize; ++it) {
            if(pointer == stackHeader[it].pointer) {
                header = stackHeader[it];
                memmove(stackHeader + it, stackHeader + it + 1, stackHeaderSize - it);
                --stackHeaderSize;
                break;
            }
        }
        if(!header.pointer) {
            return;
        }

        MemoryRange area;
        area.beginning = (U64)pointer;
        area.end = area.beginning + header.size;

        bool adjacentAreas = false;
        U64 mergedArea;
        for (U64 it = 0; it < freedStackSize; ++it) {
            if(area.beginning == freedStack[it].end) {
                if(adjacentAreas) {
                    freedStack[mergedArea].beginning = freedStack[it].beginning;
                    memmove(freedStack + it, freedStack + it + 1, freedStackSize - it);
                    --freedStackSize;
                    break;
                }

                freedStack[it].end = area.end;
                adjacentAreas = true;
                mergedArea = it;
            }
            else if(area.end == freedStack[it].beginning) {
                if(adjacentAreas) {
                    freedStack[mergedArea].end = freedStack[it].end;
                    memmove(freedStack + it, freedStack + it + 1, freedStackSize - it);
                    --freedStackSize;
                    break;
                }

                freedStack[it].beginning = area.beginning;
                adjacentAreas = true;
                mergedArea = it;
            }
        }

        if(adjacentAreas) {
            if(freedStack[mergedArea].end == stack.getMarker()) {
                StackAllocator::Marker marker = freedStack[mergedArea].beginning;
                stack.freeToMarker(marker);
                memmove(freedStack + mergedArea, freedStack + mergedArea + 1, freedStackSize - mergedArea);
                --freedStackSize;
                return;
            }
            return;
        }

        if(area.end == stack.getMarker()) {
            StackAllocator::Marker marker = area.beginning;
            stack.freeToMarker(marker);
            return;
        }

        if(freedStackSize == freedStackCapacity) {
            MemoryRange* temp = freedStack;
            freedStackCapacity *= 2;
            freedStack = new MemoryRange[freedStackCapacity];
            memcpy(freedStack, temp, freedStackSize);
            delete temp;
        }
        freedStack[freedStackSize++] = area;
    }

    void
    testMemory(void* pointer, U64 bytes, bool removeEntry) {
        if(Constants::Memory::trackMemory == 0) {
            return;
        }

        if(!removeEntry) {
            ASSERT(trackedAllocations < 655360); ///Too many allocations to track!
            for (U64 i = 0; i < trackedAllocations; ++i) {
                if(i == 184) {
                    ++--i;
                }
                ASSERT_MESSAGE(!(((U64)pointer > trackedMemory[i].beginning) && ((U64)pointer < trackedMemory[i].end)), Allocated memory was already in use!);
            }

            MemoryRange temp;
            temp.beginning = (U64)pointer;
            temp.end = temp.beginning + bytes - 1;
            trackedMemory[trackedAllocations++] = temp;
        }
        else {
            for (U64 i = 0; i < trackedAllocations; ++i) {
                if((U64)pointer == trackedMemory[i].beginning) {
                    memmove(trackedMemory + i, trackedMemory + i + 1, trackedAllocations - i);
                    --trackedAllocations;
                    return;
                }
            }
        }
    }
}

}
